#include <stdlib.h>
#include <stdio.h>
#include <errno.h>
#include <sys/resource.h>
#include <assert.h>
#include "hw1_wrappers.h"

#define NPROC_LIMIT (50)
#define SLEEP_SECS (5)

void test_fork_n()
{
    int parent_pid = getpid();
    
    printf("[U] test_fork_n\n");
    
    // test invalid params
    errno = 0;
    assert(-1 == fork_n(-1));
    assert(EINVAL == errno);
    
    assert(0 == fork_n(0));
    
    // test with allowed number of procs - should work
    //   1. create processes with fork_n, they should have incremental PIDs,
    //   2. save the generated PIDs from current_pid to youngest son pid to an array
    //   3. test the array with checkpath()

    fflush(NULL);
    if (0 != fork_n(NPROC_LIMIT / 2))
    {
        // we aren't the youngest son - drop into coma.. (for the test to work)
        printf("[U] fork_n() son pid %d\n", getpid());
        sleep(SLEEP_SECS);
        // and die!
        exit(0);
    }
        
    // we are the youngest process
    assert((parent_pid + (NPROC_LIMIT / 2)) == getpid());
    // sleep a little to make sure the other procs got to be created and sleep
    sleep(1);
    // assume pids of created sons are incremental (which is a legit assumption since
    //  we wrote fork_n(), and system calls are not preemptive in our 2.4 kernel (added in 2.6)
    {
        // + 1 size for the parent_pid
        int sons_pids[NPROC_LIMIT / 2 + 1] = {0};
        int i = 0;
        for (i = 0; i < (NPROC_LIMIT / 2 + 1); i++)
        {
            sons_pids[i] = parent_pid + i;
        }
        
        printf("[U] checking pids path from %d to %d exists\n", parent_pid, getpid());
        assert(1 == checkpath(sons_pids, NPROC_LIMIT / 2 + 1));
    }    
   
    // now test failing fork_n due to too many processes created
    
    // first make sure we aren't root, otherwise the NPROCS limit doesn't affect
    printf("[U] current effective username shouldn't be root. it is: %s\n", cuserid(NULL));
    assert(0 != strcmp("root", (char *) cuserid(NULL)));
    
    struct rlimit r = {0};
    r.rlim_cur = NPROC_LIMIT;
    r.rlim_max = NPROC_LIMIT;
    setrlimit(RLIMIT_NPROC, &r); 
    memset(&r, 0, sizeof(r));
    
    getrlimit(RLIMIT_NPROC, &r);
    printf("[U] nproc limit - hard= %d, soft= %d\n\n", (int)r.rlim_cur, (int)r.rlim_max); 
    assert((NPROC_LIMIT == r.rlim_cur) && (NPROC_LIMIT == r.rlim_max));
    
    // create sons with fork_n
    {
        int pid = getpid();
        fflush(NULL);
        assert(-1 == fork_n(NPROC_LIMIT));
        // make sure we are still the same proc
        assert(pid == getpid());
    }
    
    {
        // sleep to let sons die!
        sleep(5);
        // make sure pids don't exist
        int i = -1;
        // start from i = 2 so not to check ourselves, and our zombie son
        for (i = 2; i < NPROC_LIMIT; i++)
        {
            printf("[U] checking pid %d doesn't exist\n", getpid() + i);
            assert(-1 == kill(getpid() + i, 0));
            assert(ESRCH == errno);
        }
    }
        
    printf("[U] test_fork_n end\n");
}

#define CHECKPATH_NPROCS (10)
void test_checkpath()
{
    int sons_pids[CHECKPATH_NPROCS] = {0};
    
    printf("[U] test_checkpath\n");

    // check invalid params
    errno = 0;
    assert(-1 == checkpath(NULL, 10));
    assert(EINVAL == errno);
    
    errno = 0;
    assert(-1 == checkpath((int *)(NULL - 1), 10));
    assert(EINVAL == errno);
    
    errno = 0;
    assert(-1 == checkpath(sons_pids, 0));
    assert(EINVAL == errno);
    
    errno = 0;
    assert(-1 == checkpath(sons_pids, -10));
    assert(EINVAL == errno);
   
    
    // check special case with array of size 1 with our pid
    sons_pids[0] = getpid();
    assert(1 == checkpath(sons_pids, 1));
   
    // check array with swapper pid 0
    sons_pids[0] = 0;
    assert(1 == checkpath(sons_pids, 1));
    
    // check array with init and init, to check that init's NULL parent doesn't
    //   cuase trouble
    sons_pids[0] = 1;
    sons_pids[1] = 1;
    assert(0 == checkpath(sons_pids, 2));
 
    // check array with swapper pid, and init 1
    sons_pids[0] = 0;
    sons_pids[1] = 1;
    assert(1 == checkpath(sons_pids, 2));
    
    // check array with swapper and our pid
    sons_pids[0] = 0;
    sons_pids[1] = getpid();
    assert(0 == checkpath(sons_pids, 2));
    
    // check array with swapper, init, and our pid
    sons_pids[0] = 0;
    sons_pids[1] = 1;
    sons_pids[2] = getpid();
    assert(0 == checkpath(sons_pids, 3));
   
    // now fork to have init as parent
    // flush so sons will have empy ouput stream	
    fflush(NULL);
    if (0 == fork()) {
        sleep(1);
    } else {
        exit(0);
    }
    assert(1 == getppid());
    
    // check again    
    sons_pids[0] = 0;
    sons_pids[1] = 1;
    sons_pids[2] = getpid();
    assert(0 == checkpath(sons_pids, 3));
    
    {
        int son_num = -1;
        // create CHECKPATH_NPROCS sons, and save the pids to a local array,
        //     which will be ready for checkpath at the youngest son
        for (son_num = 0; son_num < CHECKPATH_NPROCS; son_num++) {
            int son_pid = -1;
            
            sons_pids[son_num] = getpid();
           
            // flush so sons will have empy ouput stream	
            fflush(NULL);
            // TODO assume fork succeds..
            if (0 != (son_pid = fork())) {
                printf("[U] son created %d\n", son_pid);
                break;    
            }
        }
        // the youngest son isn't in the sons_pids
        
        if (CHECKPATH_NPROCS == son_num) {
            // we are the youngest son
            int ret = -1;
            int i = -1;
            
            printf("[U] sons: ");
            for (i = 0; i < CHECKPATH_NPROCS; i++) {
                printf("%d, ", sons_pids[i]);
            }
            printf("\n");
            
            // check with all sons correct
            ret = checkpath(sons_pids, CHECKPATH_NPROCS);
            printf("[U] %d = checkpath()\n", ret);
            assert(1 == ret);
            
            // check fail with incorrect son
            {
                int correct_son;
                int j;
                
                for (j = 1; j < CHECKPATH_NPROCS; j++) {
                    correct_son = sons_pids[j];
                    // change with valid pid
                    sons_pids[j] = 1;
                    assert(0 == checkpath(sons_pids, CHECKPATH_NPROCS));
                   
                    // change with invalid pid
                    sons_pids[j] = -1;
                    assert(0 == checkpath(sons_pids, CHECKPATH_NPROCS));
     
                    sons_pids[j] = correct_son;
                }            
            }    
        }
        else {
            // we're not the youngest son - drop into coma to still exist for
		// the test of checkpath in the youngest son above
            sleep(SLEEP_SECS);
            // now die! ain't pretty.. but that's life. and who cares anyway..
            exit(0);
        }
        // the youngest son lives on, as should be.
    }
    
    printf("[U] test_checkpath end. son %d\n", getpid());
}

int main(int argc, char * argv[])
{
    printf("[U] started test\n");
    test_checkpath();
    test_fork_n();    
    printf("[U] all tests ended successful. (you can press ENTER now to see the prompt)\n");
    
    return 0;
}
